from flask import Blueprint, jsonify, request, current_app
from ..models.airport import Airport
from ..models.ammunition_depot import AmmunitionDepot
from .. import db
import math

bp = Blueprint('ammunition_supply', __name__, url_prefix='/api/ammunition-supply')

@bp.route('/initial', methods=['GET'])
def get_initial_data():
    """获取初始化数据"""
    try:
        # 获取所有区域
        regions = db.session.query(Airport.region).distinct().all()
        region_list = [r[0] for r in regions]
        
        # 获取弹药设施类型统计
        facility_types = db.session.query(
            AmmunitionDepot.type, 
            db.func.count(AmmunitionDepot.id).label('count')
        ).group_by(AmmunitionDepot.type).all()
        
        facility_stats = {t[0]: t[1] for t in facility_types}
        
        return jsonify({
            'regions': region_list,
            'facilityStats': facility_stats,
            'layouts': ['geographic', 'force', 'circular', 'hierarchical']
        })
        
    except Exception as e:
        current_app.logger.error(f'获取初始化数据失败: {str(e)}')
        return jsonify({'error': str(e)}), 500

@bp.route('/facilities', methods=['GET'])
def get_facilities():
    """获取所有弹药设施"""
    region = request.args.get('region', 'all')
    
    try:
        query = db.session.query(
            AmmunitionDepot.id,
            AmmunitionDepot.name,
            AmmunitionDepot.type,
            AmmunitionDepot.area,
            AmmunitionDepot.location,
            Airport.region,
            Airport.latitude,
            Airport.longitude
        ).join(Airport)
        
        if region != 'all':
            query = query.filter(Airport.region == region)
    
        facilities = query.all()
        
        return jsonify([{
            'id': f.id,
            'name': f.name,
            'type': f.type,
            'area': f.area,
            'location': f.location,
            'region': f.region,
            'latitude': f.latitude,
            'longitude': f.longitude
        } for f in facilities])
        
    except Exception as e:
        current_app.logger.error(f'获取弹药设施数据失败: {str(e)}')
        return jsonify({'error': str(e)}), 500

@bp.route('/network', methods=['GET'])
def get_network_data():
    """获取弹药供应网络数据"""
    region = request.args.get('region', 'all')
    
    try:
        # 获取所有弹药库数据
        query = db.session.query(
            AmmunitionDepot.id,
            AmmunitionDepot.name,
            AmmunitionDepot.type,
            AmmunitionDepot.area,
            AmmunitionDepot.location,
            Airport.region,
            Airport.latitude,
            Airport.longitude,
            Airport.name.label('airport_name')
        ).join(Airport)
        
        # 如果指定了区域，添加过滤条件
        if region != 'all':
            query = query.filter(Airport.region == region)
            
        ammo_depots = query.all()
        
        # 转换数据为前端需要的格式
        nodes = []
        links = []
        
        # 记录处理过的机场ID，用于构建连接关系
        processed_airports = {}
        
        for depot in ammo_depots:
            # 添加弹药库节点
            facility_type = get_facility_type(depot.type)
            nodes.append({
                'id': f"depot_{depot.id}",
                'name': depot.name,
                'facility_type': facility_type,
                'region': depot.region,
                'value': depot.area,  # 使用面积作为容量值
                'latitude': depot.latitude,
                'longitude': depot.longitude,
                'location': depot.location,
                'category': get_category_index(facility_type),
                'airport_id': f"airport_{depot.id}"  # 存储关联机场ID，用于建立连接
            })
            
            # 如果机场尚未处理，添加机场节点
            airport_id = f"airport_{depot.id}"
            if airport_id not in processed_airports:
                processed_airports[airport_id] = {
                    'id': airport_id,
                    'name': depot.airport_name,
                    'latitude': depot.latitude,
                    'longitude': depot.longitude,
                    'region': depot.region
                }
                
                # 添加机场节点
                nodes.append({
                    'id': airport_id,
                    'name': depot.airport_name,
                    'facility_type': 'airport',
                    'region': depot.region,
                    'value': 50,  # 默认值
                    'latitude': depot.latitude,
                    'longitude': depot.longitude,
                    'category': 3  # 机场分类
                })
                
                # 添加机场与弹药库的连接
                links.append({
                    'source': airport_id,
                    'target': f"depot_{depot.id}",
                    'value': depot.area / 5  # 连接强度，使用面积计算
                })
            
        # 为了演示网络效果，添加机场之间的连接，但仅限于同一区域内
        airports_by_region = {}
        for airport_id, airport in processed_airports.items():
            region = airport['region']
            if region not in airports_by_region:
                airports_by_region[region] = []
            airports_by_region[region].append(airport)
        
        # 同区域机场之间添加连接
        for region, airports in airports_by_region.items():
            for i in range(len(airports)):
                for j in range(i + 1, len(airports)):
                    source = airports[i]
                    target = airports[j]
                    
                    # 计算距离
                    distance = calculate_distance(
                        source['latitude'], source['longitude'],
                        target['latitude'], target['longitude']
                    )
                    
                    # 如果距离不太远，添加连接（表示物流可行性）
                    if distance < 200:  # 200公里内的机场有连接
                        links.append({
                            'source': source['id'],
                            'target': target['id'],
                            'value': 20,  # 固定值
                            'distance': distance
                        })
        
        # 定义分类
        categories = [
            {'name': '弹药制造厂'},
            {'name': '中央弹药库'},
            {'name': '机场弹药库'},
            {'name': '军用机场'}
        ]
        
        return jsonify({
            'nodes': nodes,
            'links': links,
            'categories': categories
        })
            
    except Exception as e:
        current_app.logger.error(f'获取弹药网络数据失败: {str(e)}')
        return jsonify({'error': str(e)}), 500

@bp.route('/trace/<string:facility_id>', methods=['GET'])
def trace_supply_chain(facility_id):
    """追踪弹药设施供应链"""
    current_app.logger.info(f'接收到供应链追踪请求: facility_id={facility_id}, 类型={type(facility_id)}')
    
    try:
        # 提取实际ID
        real_id = facility_id
        
        # 检查ID是否包含前缀
        if facility_id.startswith('depot_'):
            real_id = facility_id[6:]  # 去掉 "depot_" 前缀
        
        # 记录处理后的ID
        current_app.logger.info(f'处理后的ID: real_id={real_id}')
        
        # 确保ID是一个有效值
        if not real_id or real_id.isspace():
            error_msg = '无效的设施ID'
            current_app.logger.error(error_msg)
            return jsonify({'error': error_msg}), 400
        
        # 获取设施信息
        try:
            facility = db.session.query(
                AmmunitionDepot.id,
                AmmunitionDepot.name,
                AmmunitionDepot.type,
                AmmunitionDepot.area,
                AmmunitionDepot.location,
                Airport.region,
                Airport.latitude,
                Airport.longitude,
                Airport.name.label('airport_name')
            ).join(Airport).filter(AmmunitionDepot.id == real_id).first()
        except Exception as db_err:
            error_msg = f'数据库查询错误: {str(db_err)}'
            current_app.logger.error(error_msg)
            return jsonify({'error': error_msg}), 500
        
        if not facility:
            error_msg = f'未找到ID为{real_id}的设施'
            current_app.logger.error(error_msg)
            return jsonify({'error': error_msg}), 404
            
        current_app.logger.info(f'找到设施: {facility.name}, 类型: {facility.type}, 区域: {facility.region}')
        
        facility_type = get_facility_type(facility.type)
        facility_dict = {
            'id': f"depot_{facility.id}",
            'name': facility.name,
            'region': facility.region,
            'latitude': facility.latitude,
            'longitude': facility.longitude,
            'value': facility.area,
            'facility_type': facility_type,
            'location': facility.location,
            'category': get_category_index(facility_type)
        }
        
        # 添加关联机场
        airport_dict = {
            'id': f"airport_{facility.id}",
            'name': facility.airport_name,
            'region': facility.region,
            'latitude': facility.latitude,
            'longitude': facility.longitude,
            'value': 50,
            'facility_type': 'airport',
            'category': 3
        }
        
        # 获取同区域的其他设施
        related_facilities = db.session.query(
            AmmunitionDepot.id,
            AmmunitionDepot.name, 
            AmmunitionDepot.type,
            AmmunitionDepot.area,
            AmmunitionDepot.location,
            Airport.region,
            Airport.latitude,
            Airport.longitude,
            Airport.name.label('airport_name')
        ).join(Airport).filter(
            Airport.region == facility.region,
            AmmunitionDepot.id != real_id
        ).all()
        
        related_nodes = []
        related_airports = {}
        
        # 处理相关设施
        for related in related_facilities:
            related_type = get_facility_type(related.type)
            related_node = {
                'id': f"depot_{related.id}",
                'name': related.name,
                'region': related.region,
                'latitude': related.latitude,
                'longitude': related.longitude,
                'value': related.area,
                'facility_type': related_type,
                'location': related.location,
                'category': get_category_index(related_type)
            }
            related_nodes.append(related_node)
            
            # 添加相关机场
            airport_id = f"airport_{related.id}"
            if airport_id not in related_airports:
                related_airports[airport_id] = {
                    'id': airport_id,
                    'name': related.airport_name,
                    'region': related.region,
                    'latitude': related.latitude,
                    'longitude': related.longitude,
                    'value': 50,
                    'facility_type': 'airport',
                    'category': 3
                }
        
        # 构建链接
        links = []
        
        # 主设施与机场的连接
        links.append({
            'source': airport_dict['id'],
            'target': facility_dict['id'],
            'value': facility.area / 5
        })
        
        # 其他设施与各自机场的连接
        for related in related_nodes:
            related_airport_id = f"airport_{related['id'][6:]}"  # 提取数字ID
            if related_airport_id in related_airports:
                links.append({
                    'source': related_airport_id,
                    'target': related['id'],
                    'value': related['value'] / 5
                })
        
        # 添加机场之间的连接
        all_airports = [airport_dict] + list(related_airports.values())
        for i in range(len(all_airports)):
            for j in range(i + 1, len(all_airports)):
                source = all_airports[i]
                target = all_airports[j]
                
                # 计算距离
                distance = calculate_distance(
                    source['latitude'], source['longitude'],
                    target['latitude'], target['longitude']
                )
                
                # 如果距离不太远，添加连接
                if distance < 200:
                    links.append({
                        'source': source['id'],
                        'target': target['id'],
                        'value': 20,
                        'distance': distance
                    })
        
        # 合并所有节点
        all_nodes = [facility_dict, airport_dict] + related_nodes + list(related_airports.values())
        
        return jsonify({
            'facility': facility_dict,
            'nodes': all_nodes,
            'links': links
        })
        
    except Exception as e:
        current_app.logger.error(f'追踪弹药供应链失败: {str(e)}')
        return jsonify({'error': str(e)}), 500

@bp.route('/airports', methods=['GET'])
def get_airports_with_depots():
    """获取所有机场及其关联的弹药库"""
    region = request.args.get('region', 'all')
    
    try:
        # 获取所有机场数据，并关联弹药库
        query = db.session.query(
            Airport.id,
            Airport.name,
            Airport.region,
            Airport.latitude,
            Airport.longitude,
            db.func.sum(AmmunitionDepot.area).label('total_area'),
            db.func.count(AmmunitionDepot.id).label('depot_count')
        ).join(AmmunitionDepot)
        
        if region != 'all':
            query = query.filter(Airport.region == region)
            
        # 分组统计
        query = query.group_by(Airport.id)
        airports = query.all()
        
        # 为每个机场获取关联的弹药库列表
        result = []
        for airport in airports:
            # 查询该机场关联的所有弹药库
            depots = db.session.query(
                AmmunitionDepot.id,
                AmmunitionDepot.name,
                AmmunitionDepot.type,
                AmmunitionDepot.area,
                AmmunitionDepot.location
            ).filter(AmmunitionDepot.airport_id == airport.id).all()
            
            depot_list = [{
                'id': depot.id,
                'name': depot.name,
                'type': get_facility_type(depot.type),
                'area': depot.area,
                'location': depot.location
            } for depot in depots]
            
            result.append({
                'airport': {
                    'id': airport.id,
                    'name': airport.name,
                    'region': airport.region,
                    'latitude': airport.latitude,
                    'longitude': airport.longitude,
                    'total_area': airport.total_area or 0,
                    'depot_count': airport.depot_count or 0
                },
                'depots': depot_list
            })
        
        return jsonify(result)
        
    except Exception as e:
        current_app.logger.error(f'获取机场及弹药库数据失败: {str(e)}')
        return jsonify({'error': str(e)}), 500

def calculate_distance(lat1, lon1, lat2, lon2):
    """计算两个设施之间的距离（公里）"""
    # 地球半径（公里）
    R = 6371.0
    
    # 转换为弧度
    lat1 = math.radians(float(lat1))
    lon1 = math.radians(float(lon1))
    lat2 = math.radians(float(lat2))
    lon2 = math.radians(float(lon2))
    
    # 经纬度差
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    
    # Haversine公式
    a = math.sin(dlat / 2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2)**2
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
    distance = R * c
    
    return round(distance, 2)

def get_facility_type(type_code):
    """根据数据库中的type转换为前端使用的facility_type"""
    type_mapping = {
        'production': 'production',    # 弹药制造厂
        'central': 'central',          # 中央弹药库
        'airport': 'airport_depot',    # 机场弹药库
        'standard': 'standard',        # 标准弹药库
        'special': 'special'           # 特种弹药库
    }
    
    return type_mapping.get(type_code, 'standard')

def get_category_index(facility_type):
    """获取设施类型对应的分类索引"""
    category_mapping = {
        'production': 0,     # 弹药制造厂
        'central': 1,        # 中央弹药库
        'airport_depot': 2,  # 机场弹药库
        'standard': 2,       # 标准弹药库（与机场弹药库同类）
        'special': 2,        # 特种弹药库（与机场弹药库同类）
        'airport': 3         # 军用机场
    }
    
    return category_mapping.get(facility_type, 2)  # 默认为2（弹药库） 